import mindspore as ms
import mindspore.nn as nn
from util.datasets import load_mnist
from util.transform import data_iter
from util.callback import AccuracyMonitor

# from mindvision.classification.dataset import Mnist


def main():
    batch_size, lr, num_epochs = 256, 0.0001, 10

    # dataset_train = Mnist(path='/shareData/mindspore-dataset/Mnist', split="train",
    #                       batch_size=batch_size, repeat_num=1, shuffle=True,
    #                       download=False, resize=28).run()
    # dataset_test = Mnist(path='/shareData/mindspore-dataset/Mnist', split="test",
    #                      batch_size=batch_size, repeat_num=1, shuffle=True,
    #                      download=False, resize=28).run()

    # net = nn.SequentialCell(
    #     # nn.Flatten(),
    #     nn.Dense(28*28, 10),
    # )

    features_t, labels_t = load_mnist('/shareData/mindspore-dataset/Mnist/train')
    features_v, labels_v = load_mnist('/shareData/mindspore-dataset/Mnist/test', split='t10k')

    dataset_train = data_iter(features_t, labels_t, batch_size)
    dataset_valid = data_iter(features_v, labels_v, batch_size)

    net = nn.Dense(28*28, 10)

    loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    opti = nn.SGD(net.trainable_params(), learning_rate=lr)

    model = ms.Model(net, loss_fn=loss, optimizer=opti, metrics={'acc'})

    model.train(num_epochs, dataset_train, callbacks=[AccuracyMonitor(dataset_valid)])


if __name__ == '__main__':
    main()

"""
output
epoch:[1/10] Loss:2.2774138 Train Accuracy:0.14748333333333333 Valid Accuracy:0.155
epoch:[2/10] Loss:2.252627 Train Accuracy:0.19975 Valid Accuracy:0.2083
epoch:[3/10] Loss:2.2284462 Train Accuracy:0.2749666666666667 Valid Accuracy:0.2841
epoch:[4/10] Loss:2.2046874 Train Accuracy:0.3722666666666667 Valid Accuracy:0.3815
epoch:[5/10] Loss:2.1814284 Train Accuracy:0.46326666666666666 Valid Accuracy:0.4711
epoch:[6/10] Loss:2.1586423 Train Accuracy:0.5326166666666666 Valid Accuracy:0.5395
epoch:[7/10] Loss:2.136246 Train Accuracy:0.58175 Valid Accuracy:0.588
epoch:[8/10] Loss:2.11424 Train Accuracy:0.6151666666666666 Valid Accuracy:0.6216
epoch:[9/10] Loss:2.0927134 Train Accuracy:0.6402666666666667 Valid Accuracy:0.6477
epoch:[10/10] Loss:2.071551 Train Accuracy:0.6594166666666667 Valid Accuracy:0.669
"""
